package com.cs224u.experimental;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.List;

import weka.classifiers.Classifier;
import weka.classifiers.bayes.NaiveBayes;
import weka.core.Attribute;
import weka.core.FastVector;
import weka.core.Instance;
import weka.core.Instances;
import weka.filters.Filter;
import weka.filters.unsupervised.attribute.StringToWordVector;

/**
 * This is based off of the Weka sample code, which can be found here:
 * http://www.cs.waikato.ac.nz/~ml/weka/example_code/2ed/MessageClassifier.java 
 * @author ago
 *
 */
public class SimpleClassifier implements Serializable {

  private static final long serialVersionUID = 6623646459046924197L;

  //Training data
  private Instances m_Data = null;
  
  //Whether the model is up to date
  private boolean m_UpToDate;

  private StringToWordVector m_Filter = new StringToWordVector();
  
  //private Classifier m_Classifier = new J48();
  private Classifier m_Classifier = new NaiveBayes();
  
  public SimpleClassifier() throws Exception {
    String nameOfDataset = "SimpleClassifier";
    
    //TODO: Why 2?
    FastVector attributes = new FastVector(2);
    attributes.addElement(new Attribute("Message", (FastVector) null));
    
    FastVector classValues = new FastVector(2);
    classValues.addElement("normal");
    classValues.addElement("spam");
    attributes.addElement(new Attribute("Class", classValues));
    
    m_Data = new Instances(nameOfDataset, attributes, 100);
    m_Data.setClassIndex(m_Data.numAttributes() - 1); //i.e. the last column holds the label
    
    m_Classifier.setDebug(true);
    m_Filter.setLowerCaseTokens(true);
  }
  
  
	public static void main(String[] options) throws Exception {
    List<String[]> trainingAndTestData = new ArrayList<String[]>();
    //Training data
    trainingAndTestData.add(new String[]{"I like coffee", "normal"});
    trainingAndTestData.add(new String[]{"I like pizza", "normal"});
    trainingAndTestData.add(new String[]{"Do you like computers?", "normal"});
    trainingAndTestData.add(new String[]{"Winter is cold", "normal"});
    trainingAndTestData.add(new String[]{"Buy viagra", "spam"});
    trainingAndTestData.add(new String[]{"Download now", "spam"});
    trainingAndTestData.add(new String[]{"Viagra viagra viagra", "spam"});
    trainingAndTestData.add(new String[]{"Click here to try viagra", "spam"});
    trainingAndTestData.add(new String[]{"Click on this link", "spam"});
    
    //Test data
    trainingAndTestData.add(new String[]{"Viagra is good", ""});
    trainingAndTestData.add(new String[]{"I like winter", ""});
    
    SimpleClassifier messageCl = new SimpleClassifier();
    for(String[] trainingAndTestDatum : trainingAndTestData) {
      String message = trainingAndTestDatum[0];
      String classValue = trainingAndTestDatum[1];

      process(messageCl, message, classValue);
    }
	}


  private static void process(SimpleClassifier wsd, String message, String classValue)
      throws Exception {
    if(classValue.length() != 0) {
      wsd.updateData(message.toString(), classValue);
    } else {
      wsd.classifyMessage(message.toString());
    }
  }

  private void classifyMessage(String message) throws Exception {
    if(m_Data.numInstances() == 0) {
      throw new Exception("No classifier available.");
    }
    
    if(!m_UpToDate) {
      m_Filter.setInputFormat(m_Data);
      Instances filteredData = Filter.useFilter(m_Data, m_Filter);
      m_Classifier.buildClassifier(filteredData);
      m_UpToDate = true;      
      System.out.println(m_Classifier.toString());
    }
    
    Instances testset = m_Data.stringFreeStructure();
    Instance instance = makeInstance(message, testset);
    m_Filter.input(instance);
    Instance filteredInstance = m_Filter.output();
    
    for(int i = 0; i < filteredInstance.numAttributes(); i++) {
      Attribute attribute = filteredInstance.attribute(i);
      System.err.println(attribute.toString() + " " + filteredInstance.value(attribute));
    }
    
    double predicted = m_Classifier.classifyInstance(filteredInstance);
    System.err.println("Message classified as : " + m_Data.classAttribute().value((int) predicted));
  }

  private void updateData(String message, String classValue) {
    Instance instance = makeInstance(message, m_Data);
    instance.setClassValue(classValue);
    
    m_Data.add(instance);
    m_UpToDate = false;
  }

  private Instance makeInstance(String text, Instances data) {
    Instance instance = new Instance(2);
    
    Attribute messageAtt = data.attribute("Message");
    instance.setValue(messageAtt, messageAtt.addStringValue(text));
    
    instance.setDataset(data);
    return instance;
  }
}
